[YOLO-NAS Pose] Jetson AGX Orin で YOLO-NASの骨格検出を試してみました

SIN

2023.12.09

1 はじめに

CX 事業本部 delivery部の平内（SIN）です。

今年(2023年)5月に公開された YOLO-NASは、従来のYOLOv5,v8などより性能向上されたものとなっており、11月には、骨格検出モデルも追加され、オブジェクト検出モデルとして、今最も注目されているのではないでしょうか。

https://github.com/Deci-AI/super-gradients

ということで、Jetson AGX Orinで動作させて見たいと思ったのですが、ちょっと手こずったので、作業した記録として、この記事を投稿させて頂きました。

2 ImportError: cannot import name 'get_rank' from 'torch.distributed'

手こずったのは、ズバリ下記のエラーです。

# python3 -c "import super_gradients"
  File "/home/super-gradients/src/super_gradients/training/utils/distributed_training_utils.py", line 12, in <module>
    from torch.distributed import get_rank, all_gather_object
ImportError: cannot import name 'get_rank' from 'torch.distributed'

Jetson AGX Orin では、公式で準備されているモジュールでは、CUDAが利用できません。

$ pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
$ python3 -c 'import torch;print(torch.cuda.is_available())'
False

このため、通常、Nvidiaで準備されているWhlや、Dockerイメージを利用するのですが、こちらのPyTorchモジュールでは、先のエラーとなって、YOLO-NASが利用できませんでした。 https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch

エラーが出る原因は、Nvidiaで配布しているPyTorchモジュールが、USE_DISTRIBUTED=0でビルドされているためだと思います。古い配布をたどると、DISTRIBUTEDが有効になったものも有るようなのですが、ちょっとPyTorchのバージョンが古くなってしまうため、今回は、自前でビルドすることにしました。

3 PyTorchのビルド

(1) 　nvcr.io/nvidia/l4t-pytorch:r35.2.1-pth2.0-py3

Nvidiaで配布されているDockerイメージを、作業の起点としました。

$ xhost +
$ docker run -it --rm --runtime nvidia --shm-size=1g -v $(pwd)/home:/home -e DISPLAY=:0 --network host nvcr.io/nvidia/l4t-pytorch:r35.2.1-pth2.0-py3
#

(2) build

最初にビルドに必要なpyyamlをインストールしています。

# pip install pyyaml

続いて、githubからダウンロードしたソースから PyTorch をビルドしますが、この時、USE_DISTRIBUTED=1 を設定しています。

# git clone https://github.com/pytorch/pytorch.git
# cd pytorch
# export USE_DISTRIBUTED=1
# export TORCH_CUDA_ARCH_LIST="8.7"
# export PYTORCH_BUILD_VERSION=2.1.1
# export PYTORCH_BUILD_NUMBER=1
# python3 setup.py bdist_wheel
[xxxx/65xx] 小一時間かかる・・・

(3) install

ビルドが完了したら、distに作成されたwhlをインストールします

# cd dist
# pip install torch-2.1.1-cp38-cp38-linux_aarch64.whl

バージョンを確認しています。cudaや、torch.distributedが有効になっていることも確認できます。

# python3 -c "import torch;print(torch.__version__)"
2.1.1
# python3 -c 'import torch;print(torch.cuda.is_available())'
True
# python3 -c "import torch;print(torch.distributed.is_available())"
True
# python3 -c "import torch; print( torch.cuda.get_device_name(0) )"
'Orin'

(4) torchvision

YOLO-NASでは、torchvisionも必要です。PyTorch 2.1 に対応するtorchvisionは 0.16なので、これをインストールしました。

# git clone --branch release/0.16 https://github.com/pytorch/vision torchvision
# cd torchvision
export BUILD_VERSION=0.16.0
# python3 setup.py install --user

# python3 -c "import torchvision; print( torchvision.__version__ )"
0.16.0

(5) test

簡単に、動作確認しています。

# python3
Python 3.8.10 (default, Nov 14 2022, 12:59:47)
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> device = torch.device('cuda', 0)
>>> torch.rand(4).to(device)
tensor([0.2048, 0.9094, 0.4167, 0.9344], device='cuda:0')

4 super-gradientsのインストール

pip install super-gradientsでは、うまくインストールできなかったの、Githubからダウンロードしてインストールしました。

最新は、2023年11月23日リリースの3.5.0です。 https://pypi.org/project/super-gradients/#history

(1) cv2の無効化

Nvidiaで配布されたDockerイメージでは、最初からOpenCVがセットアップされているのですが、pipで管理されていないため、pip installで競合してしまいます。そこで、先にcv2を無効化しておきます。

# mv /usr/lib/python3.8/dist-packages/cv2 /usr/lib/python3.8/dist-packages/cv2.bak
# mv /usr/local/lib/python3.8/dist-packages/cv2 /usr/local/lib/python3.8/dist-packages/cv2.bak

# python3 -c 'import cv2'
ModuleNotFoundError: No module named 'cv2'

(2) setuptools -U

setuptoolsのバージョン古くてエラーが出てしまうので、最新に更新します。

# pip install setuptools -U
# pip list| grep setuptools
setuptools          69.0.2

(3) install

super-gradientsをGithubからダウンロードしてインストールします。

# git clone https://github.com/Deci-AI/super-gradients.git
# cd super-gradients
# pip install -r requirements.txt
# python3 setup.py install
# pip3 install -e .

pytorch_quantizationｍのimportでWARNINGとなっていますが、一応、利用可能になりました。

# python3 -c "import super_gradients"
The console stream is logged into /root/sg_logs/console.log
INFO - crash_tips_setup.py - Crash tips is enabled. You can set your environment variable to CRASH_HANDLER=FALSE to disable it
WARNING - __init__.py - Failed to import pytorch_quantization
INFO - font_manager.py - generated new fontManager
WARNING - calibrator.py - Failed to import pytorch_quantization
WARNING - export.py - Failed to import pytorch_quantization
WARNING - selective_quantization_utils.py - Failed to import pytorch_quantization

(4) test

画像表示のためpython3-tkをインストールします。

# apt-get install python3-tk

物体検出モデルを試してみました。

# python3
>>> import super_gradients
>>> yolo_nas = super_gradients.training.models.get("yolo_nas_l", pretrained_weights="coco").cuda()
>>> yolo_nas.predict("office.jpg", conf=0.5).show()

5 Pase(骨格検出)

セットアップが完了したので、Poseを試してみています。 yolo_nas_pose_lで、1280×720の画像を処理する時間は、0.35〜0.45sec程度でした。

import os
import time
import super_gradients

yolo_nas = super_gradients.training.models.get(
    "yolo_nas_pose_l", pretrained_weights="coco_pose"
).cuda()

input_path = "./input"
output_path = "./output"
if not os.path.exists(output_path):
    os.makedirs(output_path)

files = os.listdir(input_path)


for i, file in enumerate(files):
    if file.endswith(".jpg"):
        time_sta = time.perf_counter()

        model_predictions = yolo_nas.predict(
            "{}/{}".format(input_path, file),
            conf=0.5,
        ).save(output_path)

        time_end = time.perf_counter()
        tim = time_end - time_sta

        print("{} {} {:.2f} sec".format(i, file, tim))

        from_name = "{}/pred_0.jpg".format(output_path)
        to_name = "{}/{}".format(output_path, file)
        os.rename(from_name, to_name)

動画も試してみました。

from super_gradients.training import models
from super_gradients.common.object_names import Models

input_path = "./input_mp4"
output_path = "./output_mp4"
if not os.path.exists(output_path):
    os.makedirs(output_path)

yolo_nas_pose = models.get("yolo_nas_pose_l", pretrained_weights="coco_pose").cuda()

files = os.listdir(input_path)

for i, file in enumerate(files):
    if file.endswith(".mp4"):
        yolo_nas_pose.predict("{}/{}".format(input_path, file), conf=0.5).save(
            "{}/{}".format(output_path, file)
        )